# 导入所需的库和模块
from langchain.schema import (
    HumanMessage,
    SystemMessage
)
from langchain_openai import ChatOpenAI
import yaml
from langchain.chains import LLMChain

from langchain.memory import ConversationBufferMemory
from langchain.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

with open('config.yaml') as file:
    data = yaml.load(file, Loader=yaml.FullLoader)


# 创建一个聊天模型的实例
chatAI = ChatOpenAI(model="gpt-3.5-turbo",
                        temperature=0,
                        max_tokens=2000,
                        api_key=data['openai_api_key_github_website'],
                        base_url=data['openai_github_website_base_url'])

class CommandlineChatbot:
    def __init__(self):
        self.chatAI = chatAI
        self.messages = [SystemMessage(content="你是一个花卉行家。")]
        self.prompt = ChatPromptTemplate(
            messages=[
                SystemMessagePromptTemplate.from_template(
                    "你是一个花卉行家。你通常的回答不超过30字。"
                ),
                MessagesPlaceholder(variable_name="chat_history"),
                HumanMessagePromptTemplate.from_template("{question}")
            ]
        )
        
        # 初始化Memory
        self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
        
        # 初始化LLMChain with LLM, prompt and memory
        self.conversation = LLMChain(
            llm=self.chatAI,
            prompt=self.prompt,
            verbose=True,
            memory=self.memory
        )
    # 定义一个循环来持续与用户交互
    def chat_loop(self):
        print("Chatbot 已启动! 输入'exit'来退出程序。")
        while True:
            user_input = input("你: ")
            if user_input.lower() == 'exit':
                print("再见!")
                break
            self.messages.append(HumanMessage(content=user_input))
            response = self.conversation({"question": user_input})
            print(f"Chatbot: {response['text']}")

# 如果直接运行这个脚本，启动聊天机器人
if __name__ == "__main__":
    bot = CommandlineChatbot()
    bot.chat_loop()